#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
标签可视化脚本

功能：
1. 读取img和label文件夹中的图像和对应的mask
2. 生成可视化结果，显示原图像、mask和叠加效果
3. 支持多种颜色映射来区分不同的标签类别
4. 结果保存到同级目录的visual文件夹中

使用方法：
python utils/visualize_labels.py --img_dir /path/to/img --label_dir /path/to/label --output_dir /path/to/visual
"""

import os
import argparse
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import cv2
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors


class LabelVisualizer:
    """标签可视化器"""
    
    def __init__(self, label_colors: Dict[int, Tuple[int, int, int]] = None):
        """
        初始化标签可视化器
        
        Args:
            label_colors (Dict[int, Tuple[int, int, int]]): 标签颜色映射
        """
        # 默认颜色映射：背景=黑色，其他标签使用不同颜色
        self.label_colors = label_colors or {
            0: (0, 0, 0),        # 背景 - 黑色
            1: (0, 0, 255),      # 标签1 - 红色
            2: (0, 255, 0),      # 标签2 - 绿色
            3: (255, 0, 0),      # 标签3 - 蓝色
            4: (255, 255, 0),    # 标签4 - 黄色
            5: (255, 0, 255),    # 标签5 - 洋红色
            6: (0, 255, 255),    # 标签6 - 青色
            7: (128, 0, 128),    # 标签7 - 紫色
            8: (255, 165, 0),    # 标签8 - 橙色
            9: (128, 128, 128),  # 标签9 - 灰色
        }
        
        # 标签名称映射（可以根据实际情况修改）
        self.label_names = {
            0: "背景",
            1: "dark",
            2: "light", 
            3: "edge",
            4: "defect",
            5: "类别5",
            6: "类别6",
            7: "类别7",
            8: "类别8",
            9: "类别9",
        }
        
        print(f"标签颜色映射: {self.label_colors}")
        print(f"标签名称映射: {self.label_names}")
    
    def create_colored_mask(self, mask: np.ndarray) -> np.ndarray:
        """
        将灰度mask转换为彩色mask
        
        Args:
            mask (np.ndarray): 灰度mask图像
            
        Returns:
            np.ndarray: 彩色mask图像
        """
        height, width = mask.shape
        colored_mask = np.zeros((height, width, 3), dtype=np.uint8)
        
        # 为每个标签值分配颜色
        for label_value, color in self.label_colors.items():
            mask_region = (mask == label_value)
            colored_mask[mask_region] = color
        
        return colored_mask
    
    def create_overlay(self, image: np.ndarray, mask: np.ndarray, alpha: float = 0.5) -> np.ndarray:
        """
        创建原图像和mask的叠加效果
        
        Args:
            image (np.ndarray): 原图像
            mask (np.ndarray): 彩色mask图像
            alpha (float): 透明度，0-1之间
            
        Returns:
            np.ndarray: 叠加后的图像
        """
        # 确保图像尺寸一致
        if image.shape[:2] != mask.shape[:2]:
            mask = cv2.resize(mask, (image.shape[1], image.shape[0]))
        
        # 创建叠加效果
        overlay = cv2.addWeighted(image, 1 - alpha, mask, alpha, 0)
        return overlay
    
    def create_legend_image(self, image_size: Tuple[int, int]) -> np.ndarray:
        """
        创建图例图像
        
        Args:
            image_size (Tuple[int, int]): 图像尺寸
            
        Returns:
            np.ndarray: 图例图像
        """
        legend_height = min(200, image_size[1] // 4)
        legend_width = min(300, image_size[0] // 3)
        
        legend_img = np.ones((legend_height, legend_width, 3), dtype=np.uint8) * 255
        
        # 绘制图例
        y_offset = 20
        for label_value, color in self.label_colors.items():
            if label_value in self.label_names:
                # 绘制颜色块
                cv2.rectangle(legend_img, (10, y_offset), (30, y_offset + 15), color, -1)
                
                # 绘制标签文本
                label_text = f"{label_value}: {self.label_names[label_value]}"
                cv2.putText(legend_img, label_text, (35, y_offset + 12), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 0), 1)
                
                y_offset += 20
        
        return legend_img
    
    def visualize_single_image(self, img_path: str, label_path: str, output_path: str) -> bool:
        """
        可视化单张图像和对应的label
        
        Args:
            img_path (str): 原图像路径
            label_path (str): label图像路径
            output_path (str): 输出图像路径
            
        Returns:
            bool: 处理是否成功
        """
        try:
            # 加载原图像
            image = cv2.imread(img_path)
            if image is None:
                print(f"  ✗ 无法读取原图像: {img_path}")
                return False
            
            # 加载label图像
            label_img = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE)
            if label_img is None:
                print(f"  ✗ 无法读取label图像: {label_path}")
                return False
            
            # 创建彩色mask
            colored_mask = self.create_colored_mask(label_img)
            
            # 创建叠加效果
            overlay = self.create_overlay(image, colored_mask, alpha=0.6)
            
            # 创建图例
            legend = self.create_legend_image(image.shape[:2])
            
            # 创建可视化结果
            # 布局：原图像 | 彩色mask | 叠加效果
            h, w = image.shape[:2]
            legend_h, legend_w = legend.shape[:2]
            
            # 调整图例尺寸以匹配图像高度
            if legend_h > h:
                scale = h / legend_h
                new_w = int(legend_w * scale)
                legend = cv2.resize(legend, (new_w, h))
                legend_h, legend_w = legend.shape[:2]
            
            # 创建最终图像
            result_width = w * 3 + legend_w + 20  # 3个图像 + 图例 + 间距
            result_height = max(h, legend_h)
            result_img = np.ones((result_height, result_width, 3), dtype=np.uint8) * 255
            
            # 放置原图像
            result_img[:h, :w] = image
            
            # 放置彩色mask
            result_img[:h, w+10:w*2+10] = colored_mask
            
            # 放置叠加效果
            result_img[:h, w*2+20:w*3+20] = overlay
            
            # 放置图例 - 确保图例不会超出边界
            legend_start_x = w * 3 + 30
            legend_end_x = min(legend_start_x + legend_w, result_width)
            legend_actual_w = legend_end_x - legend_start_x
            
            if legend_actual_w > 0:
                # 调整图例宽度以匹配可用空间
                if legend_actual_w < legend_w:
                    legend = cv2.resize(legend, (legend_actual_w, legend_h))
                
                result_img[:legend_h, legend_start_x:legend_end_x] = legend
            
            # 添加标题
            titles = ["原图像", "彩色Mask", "叠加效果", "图例"]
            positions = [(w//2-50, 30), (w+w//2-50, 30), (w*2+w//2-50, 30), (w*3+legend_w//2-25, 30)]
            
            for title, (x, y) in zip(titles, positions):
                cv2.putText(result_img, title, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2)
            
            # 确保输出目录存在
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # 保存结果
            cv2.imwrite(output_path, result_img)
            
            print(f"  ✓ 保存可视化结果: {output_path}")
            return True
            
        except Exception as e:
            print(f"  ✗ 处理失败: {str(e)}")
            return False
    
    def process_directory(self, img_dir: str, label_dir: str, output_dir: str) -> Dict[str, int]:
        """
        批量处理目录中的所有图像
        
        Args:
            img_dir (str): 原图像目录
            label_dir (str): label图像目录
            output_dir (str): 输出目录
            
        Returns:
            Dict[str, int]: 处理统计信息
        """
        if not os.path.exists(img_dir):
            print(f"✗ 错误：原图像目录不存在: {img_dir}")
            return {'total': 0, 'success': 0, 'failed': 0}
        
        if not os.path.exists(label_dir):
            print(f"✗ 错误：label目录不存在: {label_dir}")
            return {'total': 0, 'success': 0, 'failed': 0}
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        print(f"✓ 输出目录: {output_dir}")
        
        # 查找所有图像文件
        img_files = []
        for file in os.listdir(img_dir):
            if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif')):
                img_files.append(file)
        
        if not img_files:
            print("✗ 警告：原图像目录中没有找到图像文件")
            return {'total': 0, 'success': 0, 'failed': 0}
        
        print(f"✓ 找到 {len(img_files)} 个图像文件")
        
        # 处理统计
        stats = {'total': len(img_files), 'success': 0, 'failed': 0}
        
        # 处理每张图像
        for i, filename in enumerate(img_files, 1):
            print(f"\n处理图像 {i}/{len(img_files)}: {filename}")
            
            # 构建文件路径
            img_path = os.path.join(img_dir, filename)
            label_filename = Path(filename).with_suffix('.png').name
            label_path = os.path.join(label_dir, label_filename)
            output_filename = f"visual_{Path(filename).stem}.jpg"
            output_path = os.path.join(output_dir, output_filename)
            
            # 检查label文件是否存在
            if not os.path.exists(label_path):
                print(f"  ⚠ label文件不存在: {label_path}")
                stats['failed'] += 1
                continue
            
            if self.visualize_single_image(img_path, label_path, output_path):
                stats['success'] += 1
            else:
                stats['failed'] += 1
        
        return stats


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='标签可视化工具')
    parser.add_argument('--img_dir', '-i', required=True, help='原图像目录路径')
    parser.add_argument('--label_dir', '-l', required=True, help='label图像目录路径')
    parser.add_argument('--output_dir', '-o', required=True, help='输出可视化结果目录路径')
    
    args = parser.parse_args()
    
    print("="*60)
    print("标签可视化工具")
    print("="*60)
    print(f"原图像目录: {args.img_dir}")
    print(f"Label目录: {args.label_dir}")
    print(f"输出目录: {args.output_dir}")
    print("-" * 60)
    
    # 创建可视化器
    visualizer = LabelVisualizer()
    
    # 批量处理
    stats = visualizer.process_directory(args.img_dir, args.label_dir, args.output_dir)
    
    # 输出统计信息
    print(f"\n{'='*60}")
    print("可视化完成！")
    print(f"总文件数: {stats['total']}")
    print(f"成功处理: {stats['success']}")
    print(f"处理失败: {stats['failed']}")
    print(f"成功率: {stats['success']/stats['total']*100:.1f}%" if stats['total'] > 0 else "成功率: 0%")
    print("="*60)


if __name__ == "__main__":
    main()
